import jwt
from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse

SECRET_KEY = "your_secret_key"  # 请使用安全的密钥

class AuthMiddleware(BaseHTTPMiddleware):
    def __init__(self, app):
        super().__init__(app)
        self.exclude_paths = ["/api/wx/callback"]  # 在类中定义排除路径

    async def dispatch(self, request: Request, call_next):
        # 检查请求路径是否在排除列表中
        if request.url.path in self.exclude_paths:
            return await call_next(request)

        # 获取认证头
        auth_header = request.headers.get('Authorization')
        if not auth_header:
            return JSONResponse(
                status_code=401,
                content={"detail": "未提供认证信息"}
            )

        try:
            # 解析 Bearer token
            scheme, token = auth_header.split()
            if scheme.lower() != 'bearer':
                raise ValueError("无效的认证方案")

            # 验证并解码 JWT
            payload = jwt.decode(token, SECRET_KEY, algorithms=["HS256"])
            
            # 将用户信息添加到请求中
            request.state.user = payload
            
            return await call_next(request)
            
        except (ValueError, jwt.InvalidTokenError):
            return JSONResponse(
                status_code=401,
                content={"detail": "无效的认证信息"}
            ) 